#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import os
import sys
from sensor_msgs.msg import PointCloud2
import open3d as o3d
import numpy as np
import struct
from datetime import datetime

class PointCloudSaver:
    def __init__(self):
        rospy.init_node('pointcloud_saver', anonymous=True)
        
        # 参数配置
        self.topic_name = rospy.get_param('~topic', '/livox/map')  # 默认话题
        self.save_path = rospy.get_param('~save_path', '/home/firefly/ws_vision/src/flypcl/pcd')  # 保存路径
        self.save_mode = rospy.get_param('~mode', 'manual')  # 保存模式: 'manual', 'auto', 'single'
        self.auto_interval = rospy.get_param('~auto_interval', 5.0)  # 自动保存间隔(秒)
        self.filename_prefix = rospy.get_param('~prefix', 'pointcloud')  # 文件名前缀
        self.file_format = rospy.get_param('~format', 'pcd')  # 文件格式: 'pcd', 'ply', 'xyz'
        
        # 创建保存目录
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
            rospy.loginfo(f"创建保存目录: {self.save_path}")
        
        # 变量初始化
        self.save_count = 0
        self.last_save_time = rospy.Time.now()
        
        # 订阅点云话题
        self.subscriber = rospy.Subscriber(self.topic_name, PointCloud2, self.pointcloud_callback)
        
        rospy.loginfo("=== 点云保存器已启动 ===")
        rospy.loginfo(f"监听话题: {self.topic_name}")
        rospy.loginfo(f"保存路径: {self.save_path}")
        rospy.loginfo(f"保存模式: {self.save_mode}")
        rospy.loginfo(f"文件格式: {self.file_format}")
        if self.save_mode == 'auto':
            rospy.loginfo(f"自动保存间隔: {self.auto_interval}秒")
        
        self.print_usage()
    
    def print_usage(self):
        """打印使用说明"""
        rospy.loginfo("=== 使用说明 ===")
        if self.save_mode == 'manual':
            rospy.loginfo("手动模式: 在终端中按 Enter 保存当前点云")
        elif self.save_mode == 'auto':
            rospy.loginfo(f"自动模式: 每 {self.auto_interval} 秒自动保存一次")
        elif self.save_mode == 'single':
            rospy.loginfo("单次模式: 接收到第一帧点云后立即保存并退出")
        rospy.loginfo("按 Ctrl+C 退出程序")
        rospy.loginfo("==================")
    
    def pointcloud_callback(self, msg):
        """点云回调函数"""
        current_time = rospy.Time.now()
        
        if self.save_mode == 'single':
            # 单次模式：保存第一帧后退出
            self.save_pointcloud(msg)
            rospy.loginfo("单次保存完成，程序退出")
            rospy.signal_shutdown("单次保存完成")
            
        elif self.save_mode == 'auto':
            # 自动模式：按时间间隔保存
            if (current_time - self.last_save_time).to_sec() >= self.auto_interval:
                self.save_pointcloud(msg)
                self.last_save_time = current_time
                
        elif self.save_mode == 'manual':
            # 手动模式：存储最新的点云，等待用户输入
            self.latest_pointcloud = msg
    
    def pointcloud2_to_o3d(self, pointcloud_msg):
        """将ROS PointCloud2消息转换为Open3D点云"""
        # 获取点云数据
        width = pointcloud_msg.width
        height = pointcloud_msg.height
        point_step = pointcloud_msg.point_step
        row_step = pointcloud_msg.row_step
        data = pointcloud_msg.data
        
        # 解析字段信息
        fields = {}
        for field in pointcloud_msg.fields:
            fields[field.name] = {
                'offset': field.offset,
                'datatype': field.datatype,
                'count': field.count
            }
        
        # 检查是否有XYZ坐标
        if 'x' not in fields or 'y' not in fields or 'z' not in fields:
            raise ValueError("点云数据缺少XYZ坐标字段")
        
        # 解析点云数据
        points = []
        colors = []
        has_color = 'rgb' in fields or ('r' in fields and 'g' in fields and 'b' in fields)
        
        for i in range(height):
            for j in range(width):
                # 计算当前点的字节偏移
                offset = i * row_step + j * point_step
                
                # 提取XYZ坐标 (假设是float32)
                x = struct.unpack_from('f', data, offset + fields['x']['offset'])[0]
                y = struct.unpack_from('f', data, offset + fields['y']['offset'])[0]
                z = struct.unpack_from('f', data, offset + fields['z']['offset'])[0]
                
                # 检查是否为有效点
                if not (np.isnan(x) or np.isnan(y) or np.isnan(z)):
                    points.append([x, y, z])
                    
                    # 提取颜色信息（如果有）
                    if has_color:
                        if 'rgb' in fields:
                            # RGB打包格式
                            rgb_data = struct.unpack_from('I', data, offset + fields['rgb']['offset'])[0]
                            r = (rgb_data >> 16) & 0xFF
                            g = (rgb_data >> 8) & 0xFF
                            b = rgb_data & 0xFF
                        elif all(field in fields for field in ['r', 'g', 'b']):
                            # 分离的RGB字段
                            r = struct.unpack_from('B', data, offset + fields['r']['offset'])[0]
                            g = struct.unpack_from('B', data, offset + fields['g']['offset'])[0]
                            b = struct.unpack_from('B', data, offset + fields['b']['offset'])[0]
                        
                        colors.append([r/255.0, g/255.0, b/255.0])
        
        # 创建Open3D点云对象
        o3d_cloud = o3d.geometry.PointCloud()
        o3d_cloud.points = o3d.utility.Vector3dVector(np.array(points))
        
        if has_color and len(colors) > 0:
            o3d_cloud.colors = o3d.utility.Vector3dVector(np.array(colors))
        
        return o3d_cloud
    
    def save_pointcloud(self, pointcloud_msg):
        """保存点云到文件"""
        try:
            # 生成文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{self.filename_prefix}_{timestamp}_{self.save_count:04d}.{self.file_format}"
            filepath = os.path.join(self.save_path, filename)
            
            # 转换PointCloud2到Open3D格式
            o3d_cloud = self.pointcloud2_to_o3d(pointcloud_msg)
            
            # 根据格式保存文件
            success = False
            if self.file_format.lower() == 'pcd':
                success = o3d.io.write_point_cloud(filepath, o3d_cloud, write_ascii=True)
            elif self.file_format.lower() == 'ply':
                success = o3d.io.write_point_cloud(filepath, o3d_cloud, write_ascii=True)
            elif self.file_format.lower() == 'xyz':
                success = o3d.io.write_point_cloud(filepath, o3d_cloud, write_ascii=True)
            else:
                rospy.logerr(f"不支持的文件格式: {self.file_format}")
                return
            
            if success:
                self.save_count += 1
                rospy.loginfo(f"✓ 已保存点云: {filename}")
                rospy.loginfo(f"  点云大小: {pointcloud_msg.width} x {pointcloud_msg.height}")
                rospy.loginfo(f"  有效点数: {len(o3d_cloud.points)}")
                rospy.loginfo(f"  文件路径: {filepath}")
                rospy.loginfo(f"  文件格式: {self.file_format.upper()}")
            else:
                rospy.logerr(f"保存点云失败: {filename}")
            
        except Exception as e:
            rospy.logerr(f"保存点云失败: {str(e)}")
    
    def manual_save_loop(self):
        """手动保存模式的主循环"""
        if self.save_mode != 'manual':
            return
            
        rospy.loginfo("手动模式已启动，按 Enter 保存点云...")
        
        while not rospy.is_shutdown():
            try:
                input()  # 等待用户按Enter
                if hasattr(self, 'latest_pointcloud'):
                    self.save_pointcloud(self.latest_pointcloud)
                else:
                    rospy.logwarn("还没有接收到点云数据，请等待...")
            except KeyboardInterrupt:
                rospy.loginfo("用户中断，程序退出")
                break
            except Exception as e:
                rospy.logerr(f"手动保存出错: {str(e)}")

def main():
    try:
        # 创建点云保存器
        saver = PointCloudSaver()
        
        # 根据模式运行
        if saver.save_mode == 'manual':
            # 手动模式需要额外的输入线程
            import threading
            
            def ros_spin():
                rospy.spin()
            
            # 启动ROS spin线程
            ros_thread = threading.Thread(target=ros_spin)
            ros_thread.daemon = True
            ros_thread.start()
            
            # 主线程处理用户输入
            saver.manual_save_loop()
            
        else:
            # 自动模式或单次模式
            rospy.spin()
            
    except rospy.ROSInterruptException:
        rospy.loginfo("程序被中断")
    except Exception as e:
        rospy.logerr(f"程序出错: {str(e)}")

if __name__ == '__main__':
    main()
